package io.gitee.luoyuan800.tool.llm.agent;

import io.gitee.luoyuan800.tool.llm.bean.AbstractMessage;
import io.gitee.luoyuan800.tool.llm.bean.Message;
import io.gitee.luoyuan800.tool.llm.bean.completion.CommonChatCompletion;
import io.gitee.luoyuan800.tool.llm.bean.config.AiAgentConfig;
import io.gitee.luoyuan800.tool.llm.bean.request.ChatRequest;
import io.gitee.luoyuan800.tool.llm.bean.request.KimiChatRequest;
import io.gitee.luoyuan800.tool.llm.web.RestfulClient;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 支持Kimi的 build in function call
 * 需要调用两次才能完成 function call
 * 这个类里面整合这段操作，所有使用的时候只需要调用一次
 * 自动完成后续的步骤返回结果
 * 目前 kimi 的build in tool 只有联网搜索
 */
@Slf4j
public class KimiAgent extends CommonAgent {
    private static final String URL = "https://api.moonshot.cn/v1/chat/completions";

    public KimiAgent(AiAgentConfig aiAgentConfig, RestfulClient restTemplate, String name) {
        super(aiAgentConfig, restTemplate, name);
        if (StringUtils.isBlank(aiAgentConfig.getUrl())) {
            aiAgentConfig.getModelConfig().setFullUrl(URL);
        }
    }

    @Override
    ChatRequest buildRequest(int resultCount) {
        KimiChatRequest request = new KimiChatRequest();
        setRequestParams(resultCount, request);
        if (getAiAgentConfig().isWebSearch()) {
            KimiChatRequest.Tool tool = new KimiChatRequest.Tool("builtin_function", "$web_search");
            request.setTools(List.of(tool));
        }
        return request;
    }

    @Override
    public CommonChatCompletion fullChat(String msg, int resultCount, String image, boolean enableThinking) {
        CommonChatCompletion commonChatCompletion = super.fullChat(msg, resultCount, image, enableThinking);
        AtomicLong token = new AtomicLong(commonChatCompletion.getUsage().getTotalTokens());
        List<CommonChatCompletion.NormalChoice> newChoices = commonChatCompletion.getChoices().stream().map(choice -> {
            boolean toolResult = "tool_calls".equalsIgnoreCase(choice.getFinishReason());
            if (toolResult) {
                log.info("Got tool call prepare message, will call again to get response...");
                Message message = choice.getMessage();
                List<AbstractMessage<?>> newMessage = new ArrayList<>(getMessages());
                newMessage.add(message);
                message.getToolCalls().forEach(toolCallResult -> {
                    Message toolMessage = new Message();
                    toolMessage.setRole(ChatRequest.MessageType.TOOL.getId());
                    toolMessage.setToolCallId(toolCallResult.getId());
                    toolMessage.setName(toolCallResult.getFunction().getName());
                    toolMessage.setContent(toolCallResult.getFunction().getArguments());
                    newMessage.add(toolMessage);
                });
                KimiChatRequest request = new KimiChatRequest();
                setRequestParams(1, request);
                request.setStream(false);
                request.setMessages(newMessage);

                CommonChatCompletion finalCommonChatCompletion = httpRequest(request);
                if (finalCommonChatCompletion != null && finalCommonChatCompletion.getChoices() != null && !finalCommonChatCompletion.getChoices().isEmpty()) {
                    token.addAndGet(finalCommonChatCompletion.getUsage().getTotalTokens());
                    log.info("Got response from {}({})", getName(), getAiAgentConfig().getModel());
                    tokenTrack(finalCommonChatCompletion);
                    return finalCommonChatCompletion.getChoices().get(0);
                } else {
                    return choice;
                }
            }
            return choice;
        }).toList();
        commonChatCompletion.setChoices(newChoices);
        commonChatCompletion.getUsage().setTotalTokens(token.intValue());
        return commonChatCompletion;
    }
}
